import torch
import matplotlib.pyplot as plt
from matplotlib.path import Path
import matplotlib.patches as patches
from mpl_toolkits.axes_grid1 import make_axes_locatable
from PPINN.trainer import PINN
import numpy as np
import math

def plot_solution(pin: PINN):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    x = torch.linspace(-1, 1, 200).to(device)
    t = torch.linspace(0, 1, 100).to(device)

    # x & t grids:
    X, T = torch.meshgrid(x, t)

    # x & t columns:
    xcol = X.reshape(-1, 1)
    tcol = T.reshape(-1, 1)

    # one large column:
    usol = pin.u_hat(xcol, tcol)

    # reshape solution:
    U = usol.reshape(x.numel(), t.numel())

    # transform to numpy:
    xnp = x.cpu().numpy()
    tnp = t.cpu().numpy()
    Unp = U.detach().cpu().numpy()

    # plot:
    fig = plt.figure(figsize=(9, 4.5))
    ax = fig.add_subplot(111)

    h = ax.imshow(Unp,
                  interpolation='nearest',
                  cmap='rainbow',
                  extent=[tnp.min(), tnp.max(), xnp.min(), xnp.max()],
                  origin='lower', aspect='auto')

    plt.xlabel('t')
    plt.ylabel('x')
    plt.title('u(x,t)', fontsize=10)

    divider = make_axes_locatable(ax)
    cax = divider.append_axes("right", size="5%", pad=0.10)
    cbar = fig.colorbar(h, cax=cax)
    cbar.ax.tick_params(labelsize=10)
    fig.savefig("test.png")

def plot_solution3D(pin: PINN, X_f, mask=None):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    coords = []
    for key, value in X_f.items():
        coords.append(X_f[key].clone().cpu().detach().requires_grad_(False))
    coords = np.vstack(coords)
    coords[:, 0] = 0
    coords = torch.tensor(np.vstack(coords), dtype=torch.float32, device=device)
    if mask is not None:
        path = Path(mask)

    # one large column:
    usol = pin.net(coords, None)
    coords = coords. cpu().detach().numpy()
    # transform to numpy:
    unp = usol.cpu().detach().numpy()
    p = unp[:, 0]
    u = np.sqrt(unp[:, 1] ** 2 + unp[:, 2] ** 2 + unp[:, 3] ** 2)
    # plot:
    fig = plt.figure(figsize=(9, 4.5))
    ax1 = fig.add_subplot(121)
    h = ax1.tricontourf(coords[:, 2], coords[:, 1], u, cmap='rainbow')
    if mask is not None:
        path = Path(mask)
        patch = patches. PathPatch(path, facecolor='none', edgecolor='k')
        ax1.add_patch(patch)
        h.set_clip_path(patch)
    plt.xlabel('z')
    plt.ylabel('y')
    plt.title('u(y,z)', fontsize=10)
    plt.axis('equal')
    divider = make_axes_locatable(ax1)
    cax = divider.append_axes("right", size="5%", pad=0.10)
    cbar = fig.colorbar(h, cax=cax)

    ax2 = fig.add_subplot(122)
    h = ax2.tricontourf(coords[:, 2], coords[:, 1], p, cmap='rainbow')
    if mask is not None:
        path = Path(mask)
        patch = patches. PathPatch(path, facecolor='none', edgecolor='k')
        ax2.add_patch(patch)
        h.set_clip_path(patch)
    plt.xlabel('z')
    plt.ylabel('y')
    plt.title('p(y,z)', fontsize=10)
    plt.axis('equal')
    divider = make_axes_locatable(ax2)
    cax = divider.append_axes("right", size="5%", pad=0.10)
    cbar = fig.colorbar(h, cax=cax)

    fig.savefig("test.png")
    
def plot_solution2D(pin: PINN, X_f):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    coords = []
    for key, value in X_f.items():
        coords.append(X_f[key])
    coords = np.vstack(coords)
    coords = torch.tensor(np.vstack(coords), dtype=torch.float32, device=device, requires_grad=True)

    # one large column:
    unp = pin.net(coords, None)
    p, ux, uy, _, _, _ = pin.u_hat(unp, coords)
    coords = coords.cpu().detach().numpy()
    p = p.cpu().detach().numpy()
    ux = ux.cpu().detach().numpy()
    uy = uy.cpu().detach().numpy()
    
    u = np.sqrt(ux ** 2 + uy ** 2)

    fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(7, 4))
    fig.subplots_adjust(hspace=0.2, wspace=0.2)
    cf = ax[0].scatter(coords[:, 0], coords[:, 1], c=u, alpha=0.4, edgecolors='none', cmap='rainbow', marker='o', s=int(3))
    ax[0].axis('square')
    for key, spine in ax[0].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    ax[0].set_xlim([min(coords[:,0]), max(coords[:,0])])
    ax[0].set_ylim([min(coords[:,1]), max(coords[:,1])])
    ax[0].set_title(r'$u$ (m/s)')
    fig.colorbar(cf, ax=ax[0], fraction=0.046, pad=0.04)

    cf = ax[1].scatter(coords[:, 0], coords[:, 1], c=p, alpha=0.4, edgecolors='none', cmap='rainbow', marker='o', s=int(3))
    ax[1].axis('square')
    for key, spine in ax[1].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[1].set_xticks([])
    ax[1].set_yticks([])
    ax[1].set_xlim([min(coords[:,0]), max(coords[:,0])])
    ax[1].set_ylim([min(coords[:,1]), max(coords[:,1])])
    ax[1].set_title('Pressure (Pa)')
    fig.colorbar(cf, ax=ax[1], fraction=0.046, pad=0.04)

    plt.savefig('./test.png', dpi=300)
    